package org.example.controller;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.ai.document.Document;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.List;

@RequiredArgsConstructor
@RestController
public class VectorController {

    private final VectorStore vectorStore;

    @GetMapping("/vec/txt/write")
    public String txt() {
        StringBuilder sb = new StringBuilder();
        ClassLoader classLoader = getClass().getClassLoader();
        InputStream inputStream = classLoader.getResourceAsStream("vectors.txt");
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
            String line;
            while ((line = reader.readLine()) != null) {
                sb.append(line);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        vectorStore.write(Arrays.stream(sb.toString().split("。")).map(Document::new).toList());
        return "ok";
    }

    @GetMapping("/vec/pdf/write")
    public String pdf() {
        vectorStore.write(this.readPdf());
        return "ok";
    }

    private List<Document> readPdf() {
        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(
                "classpath:/test.pdf",
                PdfDocumentReaderConfig
                        .builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(
                                ExtractedTextFormatter
                                        .builder()
                                        .withNumberOfTopTextLinesToDelete(0)
                                        .build()
                        )
                        .withPagesPerDocument(1)
                        .build()
        );
        return pagePdfDocumentReader.read();
    }

    @GetMapping("/vec/md/write")
    public String md() {
        vectorStore.write(this.readMd());
        return "ok";
    }

    private List<Document> readMd() {
        MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader(
                "classpath:/test.md"
        );
        return markdownDocumentReader.read();
    }
}
